import matplotlib.pyplot as plt
import matplotlib

matplotlib.rc('text', usetex=True)
matplotlib.rc('font', family='serif')

# Data for the plot
epochs = list(range(1, 11))  # Assuming we have 10 epochs
netsenseml_bandwidth = [1800, 1600, 1500, 1400, 1300, 1200, 1100, 1000, 900, 800]  # Example bandwidth values in Mbps
topk_bandwidth = [1800, 1700, 1600, 1500, 1400, 1300, 1200, 1100, 1000, 900]
allreduce_bandwidth = [1800, 1750, 1700, 1650, 1600, 1550, 1500, 1450, 1400, 1350]

# Colors for the plots
netsense_color = (221 / 255, 159 / 255, 221 / 255)
tok_color = (255 / 255, 222 / 255, 74 / 255)
allreduce_color = (112 / 255, 128 / 255, 143 / 255)

# Create the plot
plt.figure(figsize=(10, 6))
plt.plot(epochs, netsenseml_bandwidth, label='NetSenseML', color=netsense_color, marker='o')
plt.plot(epochs, topk_bandwidth, label='Top-K', color=tok_color, marker='s')
plt.plot(epochs, allreduce_bandwidth, label='AllReduce', color=allreduce_color, marker='^')

# Labeling the axes
plt.xlabel('Epoch', fontsize=20)
plt.ylabel('Training Bandwidth (samples/sec)', fontsize=20)
# plt.title('Training Bandwidth Comparison Over Epochs')

# Adding a legend
plt.legend(fontsize=20)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
# Display the plot
plt.grid(True)
plt.show()
